#!/usr/bin/env python
from racecar.data import *
from racecar.bayes import *

class Clusterator:
  pass

# Exact clustering class.
class ExactClusterator(Clusterator):
  # Requires the data object. This should be a BayesData object. 
  def __init__(self,data):
    # Keep a reference to the data.
    self._data = data;
    
    # Remember for speed's sake the number of vars and samples.
    self._nvars = data.numVars();
    self._nobs = data.numObs();
    
    # Initialize primary data structures.
    self._executed = False;
    self._done = [False]*self._nvars;
    self._finalsets = list();
    
  # Main clustering method.
  def cluster(self):
    for i in xrange(self._nvars):
      if self._done[i]:
        continue;
      else:
        tempset = set([i]);
        self._done[i] = True;
        for j in xrange(i+1,self._nvars):
          if not self._done[j]:
            if self._check(i,j):
              tempset.add(j);
              self._done[j] = True;
      self._finalsets.append(tempset);
    
    self._executed = True;
    
  # Get the number of clusters.
  def getNumClusters(self):
    if not self._executed:
      raise Exception, "You must execute the clusterator first..";
    
    return len(self._finalsets);
  
  # Get a listing of the variable clusters.
  def printClusters(self):
    if not self._executed:
      raise Exception, "You must execute a clusterator before trying to get the results! Duh!";
    else:
      count = 0;
      for varSet in self._finalsets:
        print "Cluster " + str(count) + ": " + ' '.join([self._data.getVariableName(s) for s in varSet]);
        count += 1;
        
  # Write out the clustered data.
  def writeData(self,filename,verbose=True):
    f = open(filename,'w');
    
    # Print header.
    if verbose:
      f.write("[" + " ".join([self._data.getVariableName(e) for e in self._finalsets[0]]) + "]");
      for i in xrange(1,len(self._finalsets)):
        f.write("\t" + "[" + " ".join([self._data.getVariableName(e) for e in self._finalsets[i]]) + "]");
    else:
      f.write("\t".join([str(n) for n in range(len(self._finalsets))]));
    f.write("\n");
  
    # Print the actual data.
    reps = [list(v)[0] for v in self._finalsets];
    for i in xrange(self._nobs):
      f.write("\t".join([str(self._data.getValue(rep,i)) for rep in reps]) + "\n");
    
    f.close();
    
  # Write out the cluster information.
  def writeClusterInfo(self,filename):
    if not self._executed:
      raise Exception, "You must execute a clusterator before trying to get the results! Duh!";
    else:
      f = open(filename,'w');
      
      count = 0;
      for varSet in self._finalsets:
        f.write("Cluster " + str(count) + ": " + ' '.join([self._data.getVariableName(s) for s in varSet]));
        f.write("\n");
        count += 1;
      
      f.close();
  
  # Check two variables for equivalence.
  def _check(self,var1,var2):
    forward = dict();
    reverse = dict();
    for i in xrange(self._nobs):
      v1 = self._data.getValue(var1,i);
      v2 = self._data.getValue(var2,i);
      
      # Check forward mapping from v1 to v2. 
      mvf = forward.get(v1,None);
      if mvf == None:
        forward[v1] = v2;
      else:
        if mvf != v2:
          return False;
        
      # Check reverse mapping from v2 to v1.
      mvr = reverse.get(v2,None);
      if mvr == None:
        reverse[v2] = v1;
      else:
        if mvr != v1:
          return False;
        
    return True;
  